import json
import os
import argparse

def convert_json(args):
    with open(args.input_file, 'r') as load_f:
        load_dict = json.load(load_f)

        file_list = os.listdir(args.dataset_path)
        result_list = []
        for file_name in file_list:
            if not (file_name.lower().endswith(".jpg") or file_name.lower().endswith(".jpeg")):
                continue
            image_origin_name = file_name.split('-', 1)[1]
            is_found = False
            for item in load_dict:
                if not item:
                    continue
                if image_origin_name in item.get("data").get("image"):
                    item["data"]["image"] = os.path.join(f"/labelstudio/{args.task_name}", "data/upload/1", file_name)
                    result_list.append(item)
                    is_found = True
                    break
            if not is_found:
                image_dict = dict()
                image_dict["data"] = {"image": os.path.join(f"/labelstudio/{args.task_name}", "data/upload/1", file_name)}
                image_dict['predictions'] = []
                result_list.append(image_dict)
        with open(os.path.join("convert.json"), "w") as f:
            json.dump(result_list, f)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="convert label json")
    parser.add_argument("--input_file", type=str, default='/outputs/prelabel.json')
    parser.add_argument("--dataset_path", type=str, default='/outputs/label-studio/media/upload/1')
    parser.add_argument("--task_name", type=str)
    args = parser.parse_args()
    convert_json(args)
